use rusqlite::{Connection, Result as SqliteResult, params, ToSql};
use chrono::{DateTime, Utc};
use arrow::array::{Float64Array, TimestampMillisecondArray, Array};
use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use arrow::record_batch::RecordBatch;
use std::sync::Arc;
use std::error::Error as StdError;
use anyhow::{Result, Context};

pub struct StockDataRepository {
    conn: Connection,
}

impl StockDataRepository {
    pub fn new(db_path: &str) -> SqliteResult<Self> {
        let conn = Connection::open(db_path)?;
        Ok(Self { conn })
    }

    pub fn get_historical_data(
        &self,
        stock_code: &str,
        timeframe: &str,
        start_time: Option<DateTime<Utc>>,
        end_time: Option<DateTime<Utc>>,
    ) -> Result<RecordBatch> {
        log::debug!("StockDataRepository: 获取历史数据: symbol={}, period={}, start_time={:?}, end_time={:?}",
                  stock_code, timeframe, start_time, end_time);
        let mut query = String::from(
            "SELECT timestamp, open, high, low, close, volume 
             FROM stock_candles 
             WHERE symbol = ? AND period = ?"
        );

        // 构建基本查询参数
        log::debug!("准备查询参数 - 股票代码(symbol): {}, 时间周期(period): {}", stock_code, timeframe);
        let mut query_params: Vec<&dyn ToSql> = vec![
            &stock_code,
            &timeframe,
        ];
        
        // 存储时间戳值，确保它们在查询期间保持有效
        let mut start_millis: Option<i64> = None;
        let mut end_millis: Option<i64> = None;

        if let Some(start) = start_time {
            query += " AND timestamp >= ?";
            start_millis = Some(start.timestamp_millis());
            query_params.push(start_millis.as_ref().unwrap());
        }

        if let Some(end) = end_time {
            query += " AND timestamp <= ?";
            end_millis = Some(end.timestamp_millis());
            query_params.push(end_millis.as_ref().unwrap());
        }

        query += " ORDER BY timestamp ASC";
        
        log::debug!("执行SQL查询: {}", query);
        log::debug!("查询参数数量: {}", query_params.len());

        let mut stmt = self.conn.prepare(&query)
            .context("Failed to prepare SQL statement")?;
        let rows = stmt.query_map(rusqlite::params_from_iter(query_params), |row| -> rusqlite::Result<_> {
            Ok((
                row.get::<_, i64>(0)?,
                row.get::<_, f64>(1)?,
                row.get::<_, f64>(2)?,
                row.get::<_, f64>(3)?,
                row.get::<_, f64>(4)?,
                row.get::<_, f64>(5)?,
            ))
        }).context("Failed to execute query")?;

        let mut timestamps = Vec::new();
        let mut opens = Vec::new();
        let mut highs = Vec::new();
        let mut lows = Vec::new();
        let mut closes = Vec::new();
        let mut volumes = Vec::new();

        for row in rows {
            let (ts, open, high, low, close, volume) = row.context("Failed to read row data")?;
            timestamps.push(ts);
            opens.push(open);
            highs.push(high);
            lows.push(low);
            closes.push(close);
            volumes.push(volume);
        }

        let schema = Arc::new(Schema::new(vec![
            Field::new("timestamp", DataType::Timestamp(TimeUnit::Millisecond, None), false),
            Field::new("open", DataType::Float64, false),
            Field::new("high", DataType::Float64, false),
            Field::new("low", DataType::Float64, false),
            Field::new("close", DataType::Float64, false),
            Field::new("volume", DataType::Float64, false),
        ]));

        let timestamp_array = TimestampMillisecondArray::from(timestamps);
        let open_array = Float64Array::from(opens);
        let high_array = Float64Array::from(highs);
        let low_array = Float64Array::from(lows);
        let close_array = Float64Array::from(closes);
        let volume_array = Float64Array::from(volumes);

        let columns = vec![
            Arc::new(timestamp_array) as Arc<dyn arrow::array::Array>,
            Arc::new(open_array) as Arc<dyn arrow::array::Array>,
            Arc::new(high_array) as Arc<dyn arrow::array::Array>,
            Arc::new(low_array) as Arc<dyn arrow::array::Array>,
            Arc::new(close_array) as Arc<dyn arrow::array::Array>,
            Arc::new(volume_array) as Arc<dyn arrow::array::Array>,
        ];
        
        let record_batch = RecordBatch::try_new(schema.clone(), columns)
            .context("Failed to create RecordBatch")?;
        
        log::debug!("成功创建RecordBatch，行数: {}", record_batch.num_rows());
        
        Ok(record_batch)
    }
}